
from tensorflow.keras.layers import Conv2D, MaxPooling2D, BatchNormalization, \
    Dense, Dropout, Activation, Flatten
from tensorflow.keras import backend
from tensorflow.keras.models import Sequential
from tensorflow.keras.regularizers import l1


class Mini_Vggnet:
    @staticmethod
    def build(*, width: int, height: int, depth: int, classes: int) -> Sequential:
        """
        mini_vggnet构建

        :param width: 图片宽（列）
        :param height: 图片高（行）
        :param depth: 图片通道数
        :param classes: 类别数
        :return: 模型
        """
        if backend.image_data_format() != "channels_first":
            input_shape = (height, width, depth)
            channel_dim = -1  # 记录通道维度位置
        else:
            input_shape = (depth, height, width)
            channel_dim = 1
        # 模型构建((CONV => RELU => BN) * 2 => POOL => DO)
        # 第一层
        model = Sequential()
        model.add(Conv2D(32, (3, 3), padding="same", input_shape=input_shape))
        model.add(Activation("relu"))
        model.add(BatchNormalization(axis=channel_dim))  # BN作用于通道维度

        model.add(Conv2D(32, (3, 3), padding="same"))
        model.add(Activation("relu"))
        model.add(BatchNormalization(axis=channel_dim))

        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Dropout(0.25))

        # 第二层
        model.add(Conv2D(64, (3, 3), padding="same",
                         kernel_regularizer=l1(0.01)))
        model.add(Activation("relu"))
        model.add(BatchNormalization(axis=channel_dim))

        model.add(Conv2D(64, (3, 3), padding="same"))
        model.add(Activation("relu"))
        model.add(BatchNormalization(axis=channel_dim))

        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Dropout(0.25))

        # 第三层
        model.add(Flatten())
        model.add(Dense(512, kernel_regularizer=l1(0.001)))
        model.add(Activation("relu"))
        model.add(BatchNormalization())
        model.add(Dropout(0.5))
        model.add(Dense(128, kernel_regularizer=l1(0.001)))
        model.add(Activation("relu"))
        model.add(BatchNormalization())
        model.add(Dropout(0.25))
        model.add(Dense(classes))
        model.add(Activation("softmax"))

        return model


if __name__ == '__main__':
    Mini_Vggnet.build(width=32, height=32, depth=3, classes=10)
